package scanner

import (
	"bufio"
	"fmt"
	"os"
	"strings"
	"sync"
	"kscan/pkg/config"
	"kscan/pkg/http"
	"kscan/pkg/poc"
)

// Scanner 扫描器
type Scanner struct {
	config  *config.Config
	client  *http.Client
	pocs    []*poc.POC
	targets []string
}

// NewScanner 创建扫描器
func NewScanner(cfg *config.Config) (*Scanner, error) {
	// 验证配置
	if err := cfg.Validate(); err != nil {
		return nil, err
	}

	// 创建HTTP客户端
	client, err := http.NewClient(cfg.GetTimeout(), cfg.ProxyURL)
	if err != nil {
		return nil, fmt.Errorf("创建HTTP客户端失败: %w", err)
	}

	// 加载POC
	pocs, err := poc.LoadPOCsFromDir(cfg.PocDir)
	if err != nil {
		return nil, fmt.Errorf("加载POC失败: %w", err)
	}

	if len(pocs) == 0 {
		return nil, fmt.Errorf("未找到POC文件")
	}

	// 加载目标
	targets, err := loadTargets(cfg.Target)
	if err != nil {
		return nil, fmt.Errorf("加载目标失败: %w", err)
	}

	if len(targets) == 0 {
		return nil, fmt.Errorf("未找到目标")
	}

	return &Scanner{
		config:  cfg,
		client:  client,
		pocs:    pocs,
		targets: targets,
	}, nil
}

// Scan 执行扫描
func (s *Scanner) Scan() ([]*ScanResult, error) {
	fmt.Printf("加载了 %d 个POC\n", len(s.pocs))
	fmt.Printf("加载了 %d 个目标\n", len(s.targets))
	fmt.Printf("使用 %d 个线程\n\n", s.config.Threads)

	var results []*ScanResult
	var mu sync.Mutex
	var wg sync.WaitGroup

	// 创建任务通道
	tasks := make(chan *scanTask, s.config.Threads*2)

	// 启动工作协程
	for i := 0; i < s.config.Threads; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for task := range tasks {
				result := s.scanTarget(task.target, task.poc)
				
				mu.Lock()
				results = append(results, result)
				mu.Unlock()
				
				if s.config.Verbose || result.Matched {
					s.printProgress(result)
				}
			}
		}()
	}

	// 分发任务
	go func() {
		for _, target := range s.targets {
			for _, p := range s.pocs {
				tasks <- &scanTask{
					target: target,
					poc:    p,
				}
			}
		}
		close(tasks)
	}()

	// 等待所有任务完成
	wg.Wait()

	return results, nil
}

// scanTask 扫描任务
type scanTask struct {
	target string
	poc    *poc.POC
}

// scanTarget 扫描单个目标
func (s *Scanner) scanTarget(target string, p *poc.POC) *ScanResult {
	// 执行HTTP请求
	var response *poc.Response
	var err error

	if p.Request.FollowRedirects {
		response, err = s.client.DoRequestWithRedirect(target, p.Request)
	} else {
		response, err = s.client.DoRequest(target, p.Request)
	}

	if err != nil {
		// 请求失败，返回未匹配结果
		return &ScanResult{
			Target:   target,
			PocID:    p.ID,
			PocName:  p.Name,
			Severity: p.Severity,
			Matched:  false,
		}
	}

	// 创建匹配器
	matcher := poc.NewExpressionMatcher(p.Expression)

	// 执行匹配
	matched := matcher.Match(response, p.Variables)

	// 创建结果
	return NewScanResult(target, p, matched, response)
}

// printProgress 打印进度
func (s *Scanner) printProgress(result *ScanResult) {
	if result.Matched {
		fmt.Printf("[+] 发现漏洞: %s - %s (%s)\n", result.Target, result.PocName, result.Severity)
	} else if s.config.Verbose {
		fmt.Printf("[-] 未发现: %s - %s\n", result.Target, result.PocName)
	}
}

// PrintResults 打印结果
func (s *Scanner) PrintResults(results []*ScanResult) {
	PrintResults(results, s.config.Verbose)
}

// SaveResults 保存结果
func (s *Scanner) SaveResults(results []*ScanResult, filename string) error {
	return SaveResults(results, filename)
}

// loadTargets 加载目标列表
func loadTargets(target string) ([]string, error) {
	// 检查是否是文件
	if _, err := os.Stat(target); err == nil {
		// 从文件读取目标
		return loadTargetsFromFile(target)
	}

	// 单个目标
	return []string{target}, nil
}

// loadTargetsFromFile 从文件加载目标
func loadTargetsFromFile(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var targets []string
	scanner := bufio.NewScanner(file)
	
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line != "" && !strings.HasPrefix(line, "#") {
			targets = append(targets, line)
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return targets, nil
}